import java.util.Comparator;
import java.util.HashMap;
import java.util.Map;
import java.util.PriorityQueue;

public class HuffmanTree {
    class Node{
        Node left;
        Node right;
        String code;  // 编码
        int freq;  // 频次
        Character ch; // 字符

        public Node(Character ch){
            this.ch = ch;
        }

        public Node(Node left, Node right, int freq){
            this.left = left;
            this.right = right;
            this.freq = freq;
        }

        public int getFreq(){
            return freq;
        }

        public boolean isLeaf(){
            return left == null;
        }

        @Override
        public String toString() {
            return "Node{" +
                    "left=" + left +
                    ", right=" + right +
                    ", code='" + code + '\'' +
                    ", freq=" + freq +
                    ", ch=" + ch +
                    '}';
        }
    }

    Node root;
    String str;
    Map<Character, Node> map = new HashMap<>();
    public HuffmanTree(String str){
        this.str = str;
        char[] chars = str.toCharArray();
        for (Character c : chars){
            if (!map.containsKey(c)){
                map.put(c, new Node(c));
            }
            Node node = map.get(c);
            node.freq++;
        }

        // 构建树
        PriorityQueue<Node> queue = new PriorityQueue<>(
                Comparator.comparingInt(Node::getFreq)
        );
        queue.addAll(map.values());
        while (queue.size() >= 2){
            Node x = queue.poll();
            Node y =queue.poll();
            int freq = x.getFreq() + y.getFreq();
            queue.offer(new Node(x, y, freq));
        }
        root = queue.poll();

        // 计算每个字符的编码
        // 字符串编码后占用 bits  sum为bits计算结果
        int sum = dfs(root, new StringBuilder());
        for (Node node : map.values()){
            System.out.println(node.code);
        }
    }

    private int dfs(Node node, StringBuilder stringBuilder){
        int sum = 0;
        if (node.isLeaf()){
            node.code = stringBuilder.toString();
            sum = stringBuilder.length() * node.freq;
        }else {
            sum += dfs(node.left, stringBuilder.append("0"));
            sum += dfs(node.right, stringBuilder.append("1"));
        }
        if (stringBuilder.length() > 0){
            stringBuilder.deleteCharAt(stringBuilder.length() - 1);
        }
        return sum;
    }

    // 编码
    public String enCode(){
        char[] chars = str.toCharArray();
        StringBuilder sb = new StringBuilder();
        for (char c : chars){
            sb.append(map.get(c).code);
        }
        return sb.toString();
    }

    // 解码
    public String deCode(String code){
        /*
            从根节点，寻找数字对应的字符
                数字是 0 向左走
                数字是 1 向右走
                如果没走到头，每走一步数字的索引 i++
            走到头就可以找到解码字符，再将 node 重置为根节点
         */
        char[] chars = code.toCharArray();
        int i = 0;
        StringBuilder sb = new StringBuilder();
        Node node = root;
        while (i < chars.length){
            if (!node.isLeaf()){
                if (chars[i] == '1'){
                    node = node.right;
                }else if (chars[i] == '0'){
                    node = node.left;
                }
                i++;
            }
            if (node.isLeaf()){
                sb.append(node.ch);
                node = root;
            }
        }
        return sb.toString();
    }

    public static void main(String[] args) {
        HuffmanTree tree = new HuffmanTree("abbccccccc");
        String encoded = tree.enCode();
        System.out.println(encoded);
        System.out.println(tree.deCode(encoded));
    }
}
